package ai.neuralnet.gen;

import ai.genetic.Environment;
import ai.genetic.Evolver;
import ai.genetic.Individual;
import ai.neuralnet.NNBuilder;
import ai.neuralnet.NNTrainer;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;

public class GeneticTrainer extends NNTrainer<GeneticNeuralNetwork> implements Environment<GeneticNeuralNetwork, Individual<GeneticNeuralNetwork>>
{

    protected int populationSize = 50;

    protected Evolver<GeneticNeuralNetwork, GeneticTrainer> evolver;


    protected GeneticNeuralNetwork.CrossoverStyle crossoverStyle;

    protected GeneticNeuralNetwork.MutationStyle mutationStyle;

    public GeneticTrainer(NNBuilder<?, GeneticNeuralNetwork> networkBuilder)
    {
        super(networkBuilder);
        evolver = new Evolver<GeneticNeuralNetwork, GeneticTrainer>();
        evolver.setEnvironment(this);
    }


    protected void initNetwork(GeneticNeuralNetwork network)
    {
        network.setMinGen(minRandomWeight);
        network.setMaxGen(maxRandomWeight);
        network.setCrossoverStyle(crossoverStyle);
        network.setMutationStyle(mutationStyle);
    }


    @Override
    public void trainNetworks()
    {
        networks.clear();

        int populationCount = Math.max(populationSize, Math.max(trainingCount * 2, 20));

        Collection<Individual<GeneticNeuralNetwork>> population = new ArrayList<Individual<GeneticNeuralNetwork>>(createRandomNetworks(populationCount));

        for (Individual<GeneticNeuralNetwork> p : population)
        {
            initNetwork(p.getObject());
        }

        evolver.setPopulation(population);

        List<GeneticNeuralNetwork> evolved = evolver.evolveGroup(trainingCount);

        Collections.sort(evolved);

        for (int i = 0; i < trainingCount; i++)
        {
            networks.add(evolved.get(i));
            System.out.println("Network " + (i + 1) + " trained, Final " + statsString(trainingSet.getAccuracy(networks.get(networks.size() - 1))));
        }


    }


    @Override
    public double evaluateFitness(Individual<GeneticNeuralNetwork> individual)
    {

        double f = trainingSet.getAccuracy(individual.getObject())[0];
        individual.getObject().setAccScore(f);

        return f;
    }

    @Override
    public int compare(Individual<GeneticNeuralNetwork> o1, Individual<GeneticNeuralNetwork> o2)
    {
        double diff = getFitness(o2) - getFitness(o1);
        return diff < 0 ? -1 : diff > 0 ? 1 : 0;
    }

    @Override
    public double getFitness(Individual<GeneticNeuralNetwork> individual)
    {
        return individual.getObject().getAccScore();
    }

    @Override
    public double getFitnessThreshold()
    {
        return targetAccuracy;
    }


    public void setPopulationSize(int populationSize)
    {
        this.populationSize = populationSize;
    }

    public void setBreederPercent(double breederPercent)
    {
        evolver.setBreederPercent(breederPercent);
    }


    @Override
    public void setMaxIterations(int maxIterations)
    {
        super.setMaxIterations(maxIterations);
        evolver.setMaxGenerations(maxIterations);
    }

    @Override
    public void setPrintPeriod(int printFrequency)
    {
        super.setPrintPeriod(printFrequency);
        evolver.setPrintFrequency(printFrequency);//To change body of overridden methods use File | Settings | File Templates.
    }

    @Override
    public void setPrintTrainingProgress(boolean printTrainingProgress)
    {
        super.setPrintTrainingProgress(printTrainingProgress);
        evolver.setPrintProgress(printTrainingProgress);
    }

    public void setElitism(boolean elitism)
    {
        evolver.setElitism(elitism);
    }

    public void setCrossoverStyle(GeneticNeuralNetwork.CrossoverStyle crossoverStyle)
    {
        this.crossoverStyle = crossoverStyle;
    }

    public void setMutationStyle(GeneticNeuralNetwork.MutationStyle mutationStyle)
    {
        this.mutationStyle = mutationStyle;
    }

    public void setSelectionScalar(double selectionScalar)
    {
        evolver.setSelectionScalar(selectionScalar);
    }

    public void setMortalityRate(double mortalityRate)
    {
        evolver.setMortalityRate(mortalityRate);
    }
}
